#include <libmesh/dense_matrix.h>
#include <libmesh/dense_vector.h>
#include <libmesh/dirichlet_boundaries.h>
#include <libmesh/dof_map.h>
#include <libmesh/equation_systems.h>
#include <libmesh/fe.h>
#include <libmesh/libmesh.h>
#include <libmesh/linear_implicit_system.h>
#include <libmesh/mesh.h>
#include <libmesh/mesh_generation.h>
#include <libmesh/numeric_vector.h>
#include <libmesh/quadrature_gauss.h>
#include <libmesh/sparse_matrix.h>
#include <libmesh/vtk_io.h>
#include <libmesh/zero_function.h>

#include <iostream>
#include <set>

// 函数: assemble_poisson
// 作用: 装配二维单位方形区域上的 Poisson 方程离散系统
//       -Δu = 1, 在边界上 u = 0
// 参数:
//   es:       方程系统容器
//   sys_name: 系统名称（应为 "Poisson"）
// 返回: 无
static void
assemble_poisson(libMesh::EquationSystems & es, const std::string & sys_name)
{
    using namespace libMesh;

    auto & system = es.get_system<LinearImplicitSystem>(sys_name);
    const unsigned int dim = es.get_mesh().mesh_dimension();

    const DofMap & dof_map = system.get_dof_map();
    const FEType fe_type = system.variable_type(0);

    // 构建有限元与高斯积分规则
    std::unique_ptr<FEBase> fe = FEBase::build(dim, fe_type);
    QGauss qrule(dim, SECOND);
    fe->attach_quadrature_rule(&qrule);

    const std::vector<Real> & JxW = fe->get_JxW();
    const std::vector<std::vector<Real>> & phi = fe->get_phi();
    const std::vector<std::vector<RealGradient>> & dphi = fe->get_dphi();

    DenseMatrix<Number> ke;
    DenseVector<Number> fe_local;
    std::vector<dof_id_type> dof_indices;

    for (const auto & elem : es.get_mesh().active_local_element_ptr_range())
    {
        fe->reinit(elem);
        dof_map.dof_indices(elem, dof_indices, 0);

        const unsigned int n_dofs = dof_indices.size();
        ke.resize(n_dofs, n_dofs);
        fe_local.resize(n_dofs);

        for (unsigned int qp = 0; qp < qrule.n_points(); ++qp)
        {
            for (unsigned int i = 0; i < n_dofs; ++i)
            {
                for (unsigned int j = 0; j < n_dofs; ++j)
                {
                    // 刚度项: ∫ grad(phi_i)·grad(phi_j) dΩ
                    ke(i, j) += (dphi[i][qp] * dphi[j][qp]) * JxW[qp];
                }
                // 右端项: f=1 -> ∫ phi_i dΩ
                fe_local(i) += 1.0 * phi[i][qp] * JxW[qp];
            }
        }

        // 约束并装配到全局矩阵/向量
        dof_map.constrain_element_matrix_and_vector(ke, fe_local, dof_indices);
        system.matrix->add_matrix(ke, dof_indices);
        system.rhs->add_vector(fe_local, dof_indices);
    }
}

// 函数: main
// 作用: 构建二维单位方形网格并求解 Poisson 方程，然后输出 VTK 文件进行验证
//       为避免并行测试同时写同名文件导致冲突，输出文件名包含并行规模（np）。
// 参数:
//   argc/argv: 命令行参数
// 返回: 进程退出码
int
main(int argc, char ** argv)
{
    using namespace libMesh;

    LibMeshInit init(argc, argv);
    const auto & comm = init.comm();
    const int rank = comm.rank();
    const int size = comm.size();

    if (rank == 0)
    {
        std::cout << "Poisson demo (libMesh)" << std::endl;
        std::cout << "MPI size: " << size << std::endl;
    }

    // 生成二维单位方形网格（四边形单元）
    Mesh mesh(comm);
    const unsigned int nx = 16, ny = 16;
    MeshTools::Generation::build_square(mesh, nx, ny, 0.0, 1.0, 0.0, 1.0, QUAD4);

    // 定义系统并添加变量
    EquationSystems es(mesh);
    auto & poisson = es.add_system<LinearImplicitSystem>("Poisson");
    const unsigned int u_var = poisson.add_variable("u", FIRST, LAGRANGE);

    // 为四条边界设置齐次 Dirichlet 边界条件 u=0
    std::set<boundary_id_type> bdy_ids = {0, 1, 2, 3};
    ZeroFunction<> zero_fun;
    poisson.get_dof_map().add_dirichlet_boundary(DirichletBoundary(bdy_ids, {u_var}, &zero_fun));

    // 关联装配函数并初始化与求解
    poisson.attach_assemble_function(assemble_poisson);
    es.init();
    poisson.solve();

    // 输出 VTK（并行写出 .pvtu，串行写出 .vtu），文件名包含并行规模避免测试并发冲突
    VTKIO vtk_io(mesh);
    if (size > 1)
        vtk_io.write_equation_systems(std::string("poisson_np") + std::to_string(size) + ".pvtu",
                                      es);
    else
        vtk_io.write_equation_systems(std::string("poisson_np") + std::to_string(size) + ".vtu",
                                      es);

    if (rank == 0)
        std::cout << "Poisson solve done, VTK written." << std::endl;

    return 0;
}